from torch.fx import symbolic_trace

from mymodel import *


model = MyModule()
fx_model = symbolic_trace(model)
# 打印 ATen 算子构成的图
print(fx_model.code)
# 输出示例：
# def forward(self, x, y):
#     add = torch.add(x, y);  x = y = None
#     return add

# 进一步获取底层 IR（通过 FX 的 GraphModule）
print(fx_model.graph)
# 输出：
# graph():
#   %x : [num_users=1] = placeholder[target=x]
#   %y : [num_users=1] = placeholder[target=y]
#   %add : [num_users=1] = call_function[target=torch.add](args = (%x, %y), kwargs = {})
#   return add

